import os
import re
import sys

import base

print('Step 3: Install bootloader.')

# Check
base.pacman_inst('grub efibootmgr os-prober ')


# Install
def inst_grub_bios() -> int:
    base.run('lsblk')
    print('Which disk to install grub(BIOS) bootloader?(e.g. sda)(default:none):', end='')
    disk = input()

    if not re.match('sd.', disk):
        disk = 'sda'

    # Confirm
    print(f'Install legacy bootloader in /dev/{disk}? [Y/n]', end='')
    if not base.input_yes_or_no(True):
        sys.exit()

    # Install
    return base.run(f'grub-install --target=i386-pc /dev/{disk}')


def inst_grub_uefi() -> int:
    if not os.path.exists('/boot/efi'):
        print('Mount point /boot/efi is not exist.')
        sys.exit()

    print(f'Use \'--removable\' argument? [Y/n]', end='')
    if base.input_yes_or_no(True):
        return base.run('grub-install --removable --target=x86_64-efi \
            --efi-directory=/boot/efi --bootloader-id=ARCH')
    else:
        return base.run('grub-install --target=x86_64-efi \
            --efi-directory=/boot/efi --bootloader-id=ARCH')


def grub_make_config() -> int:
    return base.run('grub-mkconfig -o /boot/grub/grub.cfg')


class Selection:
    name: str

    def __init__(self, name: str, function) -> None:
        self.name = name
        self.function = function
    
    def get_name(self) -> str:
        return self.name
    
    def inst(self) -> int:
        return self.function()


selections = (
    Selection("Grub UEFI", inst_grub_uefi),
    Selection("Grub Legacy", inst_grub_bios)
)

# Main
print()
print('These are desktop environment types selectable.')
for index in range(0, len(selections)):
    selection = selections[index]
    print(f"\t{index+1}: {selection.get_name()}")
print('Input the relevant number(s) and press Enter to install (separated by space(s)) (default:1):', end='')
opt_strs = input().split()

opts: set = set()
for opt_str in opt_strs:
    if opt_str.isdigit():
        opt = int(opt_str)-1
        if opt in range(0, len(selections)):
            opts.add(opt)
if (len(opts)) == 0:
    opts.add(0)

to_be_installed_str = ''
for opt in opts:
    to_be_installed_str += f"\'{selections[opt].get_name()}\'"

# Confirm
print()
print(f'Confirm to install boot loaders {to_be_installed_str}? [Y/n]', end='')
if not base.input_yes_or_no(True):
    sys.exit()

# Install
ret: int = 0
for opt in opts:
    if ret == 0:
        ret += selections[opt].inst()

if ret == 0:
    grub_make_config()
